import os
import json
import numpy as np
import os.path as osp
import matplotlib.pyplot as plt
# import tpdm

def load_file(file_path):
    with open(file_path, 'r') as f:
        lines = f.readlines()
    lines = [line.strip('\n') for line in lines]
    lines = [line.split(' ') for line in lines]
    return lines


def class_angle_check(class_name, lines):
    for line in lines:
        if line[0] == class_name:
            angle = line[-1]
            if abs(float(angle)) > np.pi / 2:
                return True
            else:
                return False


def angle_check(label_dir, checked_class_name, questionable_frame_out_path):
    questionable_frames = []
    files = os.listdir(label_dir)
    for file in files:
        file_path = os.path.join(label_dir, file)
        lines = load_file(file_path=file_path)
        flag = class_angle_check(checked_class_name, lines)
        if flag:
            questionable_frames.append(file.strip('.txt'))
            print(file)
    
    for frame in questionable_frames:        
        with open(questionable_frame_out_path, 'a+') as f:
            f.write(frame + '\n')
            
            
def generate_angle_distribution(root_dir, classes, out_dir, exp_name = None):
    out_dir = osp.join(out_dir, exp_name)
    if not osp.exists(out_dir):
        os.makedirs(out_dir)
    result = {
        'classes': [],
        'rotations': []
    }
    for root, _, files in os.walk(root_dir):
        for file in files:
            lines = [line.strip("\n") for line in open(osp.join(root, file), "r").readlines()]
            lines = [line.split(' ') for line in lines]
            lines = [line for line in lines if len(line) > 1]
            clses = np.array([line[0] for line in lines])
            # for _ in range(len(lines)):
            #     if lines[_][-1] == '':
            #         print(file)
            #         print(_)
            rots = np.array([float(line[-1]) for line in lines])
            result["classes"].extend(clses)
            result["rotations"].extend(rots)

    for clss in classes:
        cls_idx = np.argwhere(np.array(result["classes"]) == clss).reshape(-1)
        if len(cls_idx) > 0:
            rotations = np.array(result["rotations"])[cls_idx]
            print(clss)
            print(rotations.max())
            print(rotations.min())
            ns, bins, patches = plt.hist(rotations)
            for i in range(len(ns)):
                plt.text(bins[i], ns[i] * 1.02, int(ns[i]), fontsize=12, horizontalalignment="center")
            plt.title(f"{exp_name}_{clss} angle distribution")
            rot_max = max(rotations)
            rot_min = min(rotations)
            interval = (rot_max - rot_min) / len(ns)
            plt.xticks([rot_min + interval * i for i in range(len(ns))])
            plt.savefig(osp.join(out_dir, f"{exp_name}_{clss}.png"))
            plt.clf()
            

def check_pred_name_nums(root_dir, class_names):
    cls_nums = {k:0 for k in class_names}
    files = os.listdir(root_dir)
    for f in files:
        file_path = osp.join(root_dir, f)
        if osp.exists(file_path):
            lines = [line.strip('\n') for line in open(file_path, 'r').readlines()]
            lines = [line.split() for line in lines]
            names = [line[0] for line in lines]
            for name in names:
                if name in class_names:
                    cls_nums[name] += 1
    print(cls_nums)
    

def check_classes(dir_root, classes):
    files = os.listdir(dir_root)
    for f in files:
        file_path = os.path.join(dir_root, f)
        
        content = load_file(file_path)
        names = [i[0] for i in content]
        if classes in names:
            print(f)

if __name__ == "__main__":
    # label_dir = "/cv/yc/DSGN2/data/ww/training/label_2"
    # checked_class_name = 'lock_station'
    # questionable_frame_out_path = '/cv/yc/DSGN2/tools/check_utils/questionable_lockstation_frames.txt'
    # main(label_dir=label_dir, questionable_frame_out_path=questionable_frame_out_path, checked_class_name=checked_class_name)
    # out_dir = "/cv/yc/DSGN2/tools/check_utils/angle_distribution_imgs"
    root_dir = "/cv/yc/DSGN2/data/ww/training/label_2"
    # root_dir = "/cv/yc/DSGN2/data/kitti/training/label_2"
    classes = "tray_wo_container"
    # classes = ['Car', 'Pedestrian', 'Cyclist']
    exp_name = "ww"
    check_classes(root_dir, classes)
    # generate_angle_distribution(root_dir, classes, out_dir, exp_name)
    
    # ## check file nums
    # ori_label_dir = "/cv/yc/DSGN2/data/ww/training/label_2_withour_angle_modified"
    # new_label_dir = "/cv/yc/DSGN2/data/ww/training/label_2"
    # ori_file_num = len(os.listdir(ori_label_dir))
    # new_file_num = len(os.listdir(new_label_dir))
    # print(ori_file_num)
    # print(new_file_num)
    
    # root_dir = "/cv/yc/DSGN2/outputs/configs_stereo_kitti_models/dsgn2_wwdata_1023.stereo1023/eval/eval_with_train/epoch_60/val/final_result/data"
    
    # check_pred_name_nums(root_dir, classes)